--FEATURE-FLAG: text_to_sql

-------------------------------------------------------------------------------
-- render_semantic_catalog_sql
create or replace function ai.render_semantic_catalog_sql(id bigint, sql text, description text) returns text
as $func$
    select concat_ws
    ( E'\n'
    , format('<sql-example id=%s>', id)
    , sql
    , format('/* %s */', description)
    , '</sql-example>'
    )
$func$ language sql stable security invoker
set search_path to pg_catalog, pg_temp
;

-------------------------------------------------------------------------------
-- _render_semantic_catalog_obj_desc
create or replace function ai._render_semantic_catalog_attr_desc(id bigint, classid oid, objid oid) returns text
as $func$
declare
    _classid oid = _render_semantic_catalog_attr_desc.classid;
    _objid oid = _render_semantic_catalog_attr_desc.objid;
    _description text;
begin
    -- find descriptions of the table/view itself and columns of the table/view
    -- the table/view itself has an objsubid = 0
    -- columns have objsubid > 0
    -- prefer descriptions from the semantic catalog, fallback to COMMENTs
    with x as
    (
        select
          f.attname -- null for the table/view itself
        , coalesce(a.objsubid, b.objsubid) as objsubid
        , coalesce(a.description, b.description) as description
        from
        (
            select x.objsubid, x.description
            from ai.semantic_catalog_obj x
            where x.classid = _classid
            and x.objid = _objid
        ) a
        full outer join
        (
            select d.objsubid, d.description
            from pg_catalog.pg_description d
            where d.classoid = _classid
            and d.objoid = _objid
        ) b
        on (a.objsubid = b.objsubid)
        left outer join pg_catalog.pg_attribute f -- left outer b/c the table/view itself won't have a row in pg_attribute
        on (f.attrelid = _objid and f.attnum = coalesce(a.objsubid, b.objsubid))
    )
    select 
      string_agg
      (
        case
            when x.objsubid = 0 and k.relkind in ('r', 't', 'f', 'p', 'v', 'm') then
                format
                ( $sql$COMMENT ON %s %I.%I IS $$%s$$;$sql$
                , case k.relkind 
                    when 'f' then 'FOREIGN TABLE'
                    when 'v' then 'VIEW' 
                    when 'm' then 'MATERIALIZED VIEW'
                    else 'TABLE'
                  end
                , n.nspname
                , k.relname
                , x.description
                )
            else
                format($sql$COMMENT ON COLUMN %I.%I.%I IS $$%s$$;$sql$, n.nspname, k.relname, x.attname, x.description)
        end
      , E'\n' order by x.objsubid nulls first
      )
    into _description
    from x
    inner join pg_class k on (k.oid = _objid)
    inner join pg_catalog.pg_namespace n on (k.relnamespace = n.oid)
    where x.description is not null -- shouldn't be the case but paranoia strikes deep
    ;
    return _description;
end
$func$ language plpgsql stable security invoker
set search_path to pg_catalog, pg_temp
;

-------------------------------------------------------------------------------
-- _render_semantic_catalog_table
create or replace function ai._render_semantic_catalog_table(id bigint, classid oid, objid oid) returns text
as $func$
declare
    _classid oid = _render_semantic_catalog_table.classid;
    _objid oid = _render_semantic_catalog_table.objid;
    _nspname pg_catalog.name;
    _relname pg_catalog.name;
    _persistence pg_catalog.text;
    _columns pg_catalog.text[];
    _constraints pg_catalog.text[];
    _indexes pg_catalog.text;
    _ddl text;
    _description text;
    _sample_data text;
begin
    -- get the descriptions (if any) of the table and the columns of the table
    select ai._render_semantic_catalog_attr_desc(id, _classid, _objid)
    into _description
    ;

    -- names
    select
      n.nspname
    , k.relname
    , case k.relpersistence
        when 't' then 'TEMPORARY'
        when 'u' then 'UNLOGGED'
        else ''
      end
    into strict
      _nspname
    , _relname
    , _persistence
    from pg_catalog.pg_class k
    inner join pg_catalog.pg_namespace n
    on (k.relnamespace operator(pg_catalog.=) n.oid)
    where k.oid operator(pg_catalog.=) objid
    ;

    -- columns
    select pg_catalog.array_agg(x.txt order by x.attnum)
    into strict _columns
    from
    (
        select pg_catalog.concat_ws
        ( ' '
        , a.attname
        , pg_catalog.format_type(a.atttypid, a.atttypmod)
        , case when a.attnotnull then 'NOT NULL' else '' end
        , case
            when a.atthasdef
                then pg_catalog.pg_get_expr(d.adbin, d.adrelid)
            when a.attidentity operator(pg_catalog.=) 'd'
                then 'GENERATED BY DEFAULT AS IDENTITY'
            when a.attidentity operator(pg_catalog.=) 'a'
                then 'GENERATED ALWAYS AS IDENTITY'
            when a.attgenerated operator(pg_catalog.=) 's'
                then pg_catalog.format('GENERATED ALWAYS AS (%s) STORED', pg_catalog.pg_get_expr(d.adbin, d.adrelid))
            else ''
          end
        ) as txt
        , a.attnum
        from pg_catalog.pg_attribute a
        left outer join pg_catalog.pg_attrdef d
        on (a.attrelid operator(pg_catalog.=) d.adrelid and a.attnum operator(pg_catalog.=) d.adnum)
        where a.attrelid operator(pg_catalog.=) objid
        and a.attnum operator(pg_catalog.>) 0
        and not a.attisdropped
    ) x;

    -- constraints
    select pg_catalog.array_agg(pg_catalog.pg_get_constraintdef(k.oid, true) order by k.conname)
    into _constraints
    from pg_catalog.pg_constraint k
    where k.conrelid operator(pg_catalog.=) objid
    ;

    -- indexes
    select coalesce(pg_catalog.string_agg(pg_catalog.pg_get_indexdef(i.indexrelid, 0, true), E';\n'), '')
    into strict _indexes
    from pg_catalog.pg_index i
    where i.indrelid operator(pg_catalog.=) objid
    ;

    -- ddl
    select pg_catalog.format(E'CREATE %s TABLE %I.%I\n( ', _persistence, _nspname, _relname)
    operator(pg_catalog.||)
    pg_catalog.string_agg(x.line, E'\n, ')
    operator(pg_catalog.||)
    E'\n);\n'
    operator(pg_catalog.||) coalesce(_indexes, '')
    into strict _ddl
    from
    (
        select * from pg_catalog.unnest(_columns) line
        union all
        select * from pg_catalog.unnest(_constraints) line
    ) x
    ;

    if _ddl is null then
        return '';
    end if;
    
    select ai._render_sample(pg_catalog.format('%I.%I', _nspname, _relname))
    into _sample_data
    ;

    return concat_ws
    ( E'\n'
    , format('<table id=%s>', id)
    , _ddl
    , coalesce(_description, '')
    , case when _sample_data is not null then
        E'-- sample data\n' || _sample_data
      else ''
      end
    , '</table>'
    );
end
$func$ language plpgsql stable security invoker
set search_path to pg_catalog, pg_temp
;

-------------------------------------------------------------------------------
-- _render_semantic_catalog_view
create or replace function ai._render_semantic_catalog_view(id bigint, classid oid, objid oid) returns text
as $func$
declare
    _classid pg_catalog.oid = _render_semantic_catalog_view.classid;
    _objid pg_catalog.oid = _render_semantic_catalog_view.objid;
    _view_schema pg_catalog.name;
    _view_name pg_catalog.name;
    _ddl pg_catalog.text;
    _description pg_catalog.text;
    _sample_data pg_catalog.text;
begin
    -- get the descriptions (if any) of the view and the columns of the view
    select ai._render_semantic_catalog_attr_desc(id, _classid, _objid)
    into _description
    ;

    -- get the ddl for the view
    perform 1
    from pg_extension
    where extname = 'timescaledb'
    ;
    if found then
        select n.nspname, c.relname
        into _view_schema, _view_name
        from pg_class c
        join pg_namespace n on n.oid = c.relnamespace
        where c.oid = _objid;

        select view_definition
        into _ddl
        from timescaledb_information.continuous_aggregates
        where view_schema = _view_schema and view_name = _view_name
        ;
    end if;

    if _ddl is null then
        select pg_catalog.pg_get_viewdef(_objid, true)
        into _ddl
        ;
    end if;

    if _ddl is null then
        return '';
    end if;

    select ai._render_sample(pg_catalog.format('%I.%I', n.nspname, k.relname))
    into _sample_data
    from pg_catalog.pg_class k
    inner join pg_catalog.pg_namespace n on (k.relnamespace = n.oid)
    where k.oid = _objid
    ;

    return concat_ws
    ( E'\n'
    , format('<view id=%s>', id)
    , _ddl
    , coalesce(_description, '')
    , case when _sample_data is not null then
        E'-- sample data\n' || _sample_data
      else ''
      end
    , '</view>'
    );
end
$func$ language plpgsql stable security invoker
set search_path to pg_catalog, pg_temp
;

-------------------------------------------------------------------------------
-- _render_semantic_catalog_func
create or replace function ai._render_semantic_catalog_func(id bigint, classid oid, objid oid) returns text
as $func$
declare
    _classid oid = _render_semantic_catalog_func.classid;
    _objid oid = _render_semantic_catalog_func.objid;
    _nspname name;
    _proname name;
    _prokind char;
    _idargs text;
    _ddl text;
    _description text;
begin

    select
      n.nspname
    , p.proname
    , case p.prokind
        when 'f' then 'function'
        when 'w' then 'function'
        when 'p' then 'procedure'
        when 'a' then 'aggregate'
      end
    into
      _nspname
    , _proname
    , _prokind
    from pg_catalog.pg_proc p
    inner join pg_catalog.pg_namespace n on (p.pronamespace = n.oid)
    where p.oid = _objid
    ;
    
    select pg_get_function_identity_arguments(_objid) into _idargs;

    select x.description into _description
    from ai.semantic_catalog_obj x
    where x.classid = _classid
    and x.objid = _objid
    and x.objsubid = 0
    ;
    if _description is null then
        select d.description into _description
        from pg_catalog.pg_description d
        where d.classoid = _classid
        and d.objoid = _objid
        and d.objsubid = 0
        ;
    end if;

    select pg_catalog.pg_get_functiondef(_objid) into _ddl;

    if _ddl is null then
        return '';
    end if;

    return concat_ws
    ( E'\n'
    , format
      ( '<%s id=%s>'
      , _prokind
      , id
      )
    , _ddl
    , case when _description is null then '' else
        format
        ( $sql$COMMENT ON %s %I.%I(%s) IS $$%s$$;$sql$
        , _prokind
        , _nspname
        , _proname
        , _idargs
        , _description
        )
      end
    , format('</%s>', _prokind)
    );
end
$func$ language plpgsql stable security invoker
set search_path to pg_catalog, pg_temp
;

-------------------------------------------------------------------------------
-- render_semantic_catalog_obj
create or replace function ai.render_semantic_catalog_obj(id bigint, classid oid, objid oid) returns text
as $func$
    select
        case k.relname
            when 'pg_class' then
                case
                    when k.relkind in ('r', 'p') then
                        ai._render_semantic_catalog_table(id, classid, objid)
                    when k.relkind in ('v', 'm') then
                        ai._render_semantic_catalog_view(id, classid, objid)
                    else ''
                    end
            when 'pg_proc' then
                ai._render_semantic_catalog_func(id, classid, objid)
            else ''
        end
    from pg_catalog.pg_class k
    where k.oid = classid
$func$ language sql stable security invoker
set search_path to pg_catalog, pg_temp
;

-------------------------------------------------------------------------------
-- render_semantic_catalog_obj_listing
create or replace function ai.render_semantic_catalog_obj_listing(n pg_catalog.int8) returns text
as $func$
    select concat_ws
    ( E'\n'
    , 'Below is a list of commonly used database objects which may or may not be relevant.'
    , '<commonly-used-database-objects>'
    , string_agg(x.item, E'\n')
    , E'</commonly-used-database-objects>\n'
    )
    from
    (
        select
          case o.objtype
            when 'table' then
                format('%s. table: %I.%I', o.id, o.objnames[1], o.objnames[2])
            when 'view' then
                format('%s. view: %I.%I', o.id, o.objnames[1], o.objnames[2])
            when 'function' then
                format('%s. function: %s', o.id, o.objid::regprocedure)
          end as item
        from ai.semantic_catalog_obj o
        where o.objsubid = 0
        order by o.usage desc
        limit n
    ) x
    ;
$func$ language sql stable security invoker
set search_path to pg_catalog, pg_temp
;

-------------------------------------------------------------------------------
-- _resolve_class
create or replace function ai._resolve_class
( ambiguous pg_catalog.text
, search_path pg_catalog.text default pg_catalog.current_setting('search_path')
, out fully_qualified pg_catalog.text -- fully qualified and properly quoted if necessary
, out schema pg_catalog.name -- not quoted
, out object pg_catalog.name -- not quoted
)
as $func$
declare
    _parts pg_catalog.text[];
begin
    _parts = pg_catalog.parse_ident(ambiguous, false);
    assert pg_catalog.array_length(_parts, 1) in (1, 2);
    
    -- if it's already schema qualified, just convert to regclass
    if pg_catalog.array_length(_parts, 1) = 2 then
        schema = _parts[1];
        object = _parts[2];
        fully_qualified = pg_catalog.concat_ws
        ( '.'
        , pg_catalog.quote_ident(schema)
        , pg_catalog.quote_ident(object)
        );
        -- it's fully qualified, but does it exist?
        if pg_catalog.to_regclass(fully_qualified) is null then
            fully_qualified = null;
            schema = null;
            object = null;
            return;
        end if;
        return;
    end if;
    
    -- check each schema in the search_path in order to find the thing
    object = _parts[1];
    for schema in
    (
        select 
          case when x.part = '"$user"' then pg_catalog."current_user"() 
          else (pg_catalog.parse_ident(trim(x.part)))[1] -- remove double quotes if exists
          end as schema
        from pg_catalog.string_to_table(search_path, ',') x(part)
        union
        values -- implicit search_path
          ('pg_catalog')
        , ('pg_temp')
    )
    loop
        fully_qualified = pg_catalog.concat_ws
        ( '.'
        , pg_catalog.quote_ident(schema)
        , pg_catalog.quote_ident(object)
        );
        if pg_catalog.to_regclass(fully_qualified) is not null then
            -- hooray! we found it!
            return;
        end if;
    end loop;

    -- we didn't find it :(
    fully_qualified = null;
    schema = null;
    object = null;
    return;
end
$func$ language plpgsql stable security invoker
set search_path to pg_catalog, pg_temp
;

-------------------------------------------------------------------------------
-- _render_obj_sample
create or replace function ai._render_sample
( fully_qualified_relation pg_catalog.text
, total pg_catalog.int4 default 3
) returns text
as $python$
    #ADD-PYTHON-LIB-DIR
    import ai.semantic_catalog
    return ai.semantic_catalog.render_sample(plpy, fully_qualified_relation, total)
$python$
language plpython3u volatile parallel safe security invoker
set search_path to pg_catalog, pg_temp
;

-------------------------------------------------------------------------------
-- render_obj_sample
create or replace function ai.render_sample
( relation pg_catalog.text
, total pg_catalog.int4 default 3
, search_path pg_catalog.text default pg_catalog.current_setting('search_path', true)
) returns text
as $func$
declare
    _fully_qualified pg_catalog.text;
begin
    select fully_qualified into _fully_qualified
    from ai._resolve_class(relation, search_path)
    ;
    if _fully_qualified is null then
        return null;
    end if;
    return ai._render_sample(_fully_qualified, total);
end
$func$
language plpgsql volatile parallel safe security invoker
set search_path to pg_catalog, pg_temp
;
